import operator
from abc import ABC, abstractmethod
from functools import reduce
from typing import Any, Dict, List, Union

from ..StrategiesVector import StrategiesVector
from ..shard.placement_types import OperationData,DeviceMesh,ShardingStrategy



class StrategyGenerator(ABC):
    """
    StrategyGenerator is used to generate the same group of sharding strategies.

    TODO: remove the original strategy_generator.py after refactoring
    """

    def __init__(self, operation_data_mapping: Dict[str, OperationData], device_mesh: DeviceMesh):
        self.op_data = operation_data_mapping   #用于后续的cost计算
        self.device_mesh = device_mesh
        # self.node_name = node_name

    def generate(self) -> List[ShardingStrategy]:
        """
        Generate all possible sharding strategies for this operation.
        """
        # 枚举所有可能的策略
        strategies = self.collate_strategies()

        # some strategies may be None as ignore_sharding_exception may return None
        # when ShardingSpecException occurs.
        # thus, remove those None values
        strategies = [strategy for strategy in strategies if strategy]

        # update the costs
        # update mete info on cost
        # these update methods are all in-place, the default method will do nothing
        # the cost info will only be added if the child class overrides these methods
        # for strategy in strategies:
            # self.update_communication_cost(strategy)
            # self.update_compute_cost(strategy)
            # self.update_memory_cost(strategy)

        return strategies

    @abstractmethod
    def collate_strategies(self) -> List[ShardingStrategy]:
        pass

    # @abstractmethod
    # def validate(self) -> bool:
    #     """
    #     Validate if the operands are of desired shape.
    #     If True, means this generator can be used for the current operation.
    #     """
    #     pass


class FollowingStrategyGenerator(StrategyGenerator):
    """
    FollowingStrategyGenerator is used to generate the sharding strategies which depends on its predecessor node.

    TODO: remove the original strategy_generator.py after refactoring
    """

    def __init__(self, operation_data_mapping: Dict[str, OperationData], device_mesh: DeviceMesh,pre_strategies_vectors: Dict[str,List[StrategiesVector]]):
        self.op_data = operation_data_mapping
        self.device_mesh = device_mesh
        self.pre_strategies_vectors = pre_strategies_vectors
        # 不以最少的，以第一个
        for index, node_name in enumerate(pre_strategies_vectors):
            if index == 0:
                self.pre_strategies_vector = pre_strategies_vectors[node_name]
                break

